package lock

import (
	"fmt"
	"sync"
	"time"
)

type MyRWMutex struct {
	mu        sync.Mutex
	readers   int
	writing   bool
	readCond  *sync.Cond
	writeCond *sync.Cond
}

// NewMyRWMutex 初始化自定义读写锁
func NewMyRWMutex() *MyRWMutex {
	m := &MyRWMutex{}
	m.readCond = sync.NewCond(&m.mu)
	m.writeCond = sync.NewCond(&m.mu)
	return m
}

// RLock 加读锁
func (m *MyRWMutex) RLock() {
	m.mu.Lock()
	// 若有写操作正在进行，读操作等待
	for m.writing {
		m.readCond.Wait()
	}
	m.readers++
	m.mu.Unlock()
}

// RUnlock 解读锁
func (m *MyRWMutex) RUnlock() {
	m.mu.Lock()
	m.readers--
	// 若没有读操作了，唤醒可能等待的写操作
	if m.readers == 0 {
		m.writeCond.Signal()
	}
	m.mu.Unlock()
}

// Lock 加写锁
func (m *MyRWMutex) Lock() {
	m.mu.Lock()
	// 若有其他写操作或读操作正在进行，写操作等待
	for m.writing || m.readers > 0 {
		m.writeCond.Wait()
	}
	m.writing = true
	m.mu.Unlock()
}

// Unlock 解写锁
func (m *MyRWMutex) Unlock() {
	m.mu.Lock()
	m.writing = false
	// 唤醒所有等待的读操作和一个等待的写操作
	m.readCond.Broadcast()
	m.writeCond.Signal()
	m.mu.Unlock()
}

// 共享资源
var sharedData int

// 自定义读写锁实例
var myRWMutex = NewMyRWMutex()

// 读操作函数
func readData(id int) {
	myRWMutex.RLock()
	defer myRWMutex.RUnlock()
	fmt.Printf("Reader %d reads data: %d\n", id, sharedData)
	time.Sleep(100 * time.Millisecond)
}

// 写操作函数
func writeData(id int) {
	myRWMutex.Lock()
	defer myRWMutex.Unlock()
	sharedData++
	fmt.Printf("Writer %d writes data: %d\n", id, sharedData)
	time.Sleep(200 * time.Millisecond)
}

func TryRWLock() {
	var wg sync.WaitGroup

	// 启动多个读协程
	for i := 0; i < 3; i++ {
		wg.Add(1)
		go func(id int) {
			defer wg.Done()
			readData(id)
		}(i)
	}

	// 启动写协程
	wg.Add(1)
	go func(id int) {
		defer wg.Done()
		writeData(id)
	}(0)

	// 等待所有协程执行完毕
	wg.Wait()
	fmt.Println("All operations are done.")
}
